package com.crawler.search.utils;

import cn.hutool.core.exceptions.ExceptionUtil;
import cn.hutool.core.util.ArrayUtil;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.crawler.search.constant.CommonColumn;
import com.crawler.search.constant.ESType;
import com.crawler.search.constant.QuerierMapper;
import com.crawler.search.constant.ResultMapper;
import com.crawler.search.index.entity.QueryParam;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
import org.elasticsearch.search.rescore.QueryRescoreMode;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.sort.ScoreSortBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 基本查询
 *
 * @author zyh
 */
@Component
public class QueryParseUtil {

    protected static Logger log = LoggerFactory.getLogger(QueryParseUtil.class);

    @Autowired
    protected RestHighLevelClient client;

    /**
     * 查询拼接 SearchRequestBuilder
     *
     * @param param
     * @return
     */
    public SearchSourceBuilder buildSearch(QueryParam param) {

        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(param.getQuery());

        buildHighlight(param, searchSourceBuilder);//高亮
        buildRescore(param, searchSourceBuilder);//重打分
        buildSort(param, searchSourceBuilder);//排序
        buildFetchSource(param, searchSourceBuilder);//返回字段设置

        JSONObject obj = JSONObject.parseObject( param.getSearch());
        int pageNo = obj.containsKey(QuerierMapper.PAGE_NO) ? obj.getIntValue(QuerierMapper.PAGE_NO) : CommonColumn.PAGENO;
        int pageSize = obj.containsKey(QuerierMapper.PAGE_SIZE) ? obj.getIntValue(QuerierMapper.PAGE_SIZE) : CommonColumn.PAGESIZE;

        pageSize = pageSize == 0 ? CommonColumn.PAGESIZE : pageSize;
        searchSourceBuilder.from((pageNo - CommonColumn.PAGENO) * pageSize)
                .size(pageSize);//分页
        return searchSourceBuilder;
    }

    /**
     * 查询
     *
     * @param
     * @return
     */
    public String getSearchResult(QueryParam param) {
        SearchSourceBuilder srb = buildSearch(param);
        log.info("[es-query] {}-\n index:{} type:{} query:\n{} ", "getSearchResult", param.getSearchIndex(), param.getSearchType().split(","), srb);
        long begin = System.currentTimeMillis();
        SearchRequest searchRequest = new SearchRequest(param.getSearchIndex())
                .searchType(SearchType.DEFAULT)
                .source(srb);
        SearchResponse searchResponse = null;
        try {
            searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("[es-query] error {} ", ExceptionUtil.stacktraceToString(e));
        }
        log.info("[es-query] {}- 耗时：{} ms", "getSearchResult", (System.currentTimeMillis() - begin));
        JSONObject obj = JSONObject.parseObject( param.getSearch());
        int pageNo = obj.containsKey(QuerierMapper.PAGE_NO) ? obj.getIntValue(QuerierMapper.PAGE_NO) : 1;
        int pageSize = obj.containsKey(QuerierMapper.PAGE_SIZE) ? obj.getIntValue(QuerierMapper.PAGE_SIZE) : 10;
        return handleResponse(searchResponse, pageNo, pageSize);
    }

    /**
     * 查询结果处理
     *
     * @param sr
     * @return
     */
    public static String handleResponse(SearchResponse sr, int pageNo, int pageSize) {
        JSONObject response = new JSONObject();
        response.put(ResultMapper.COUNT, sr.getHits().getTotalHits().value);
        response.put(ResultMapper.PAGENO, pageNo);
        response.put(ResultMapper.PAGESIZE, pageSize);
        SearchHits hits = sr.getHits();
        SearchHit sh = null;
        List<JSONObject> list = new ArrayList<JSONObject>();
        for (int i = 0; i < hits.getHits().length; i++) {
            try {
                sh = hits.getAt(i);
                String index = sh.getIndex();
                JSONObject hitJO = JSONObject.parseObject(sh.getSourceAsString());
                hitJO.put(ResultMapper.INDEX, index);
                Map<String, HighlightField> highlightFields = sh.getHighlightFields();
                if (highlightFields != null) {
                    for (Map.Entry<String, HighlightField> entry : highlightFields.entrySet()) {
                        Text[] texts = entry.getValue().getFragments();
                        hitJO.put(entry.getKey(), texts[0].toString());
                    }
                }
                list.add(hitJO);
            } catch (Exception e) {
            }
        }
        JSONArray arr = JSONArray.parseArray(JSONObject.toJSONString(list));
        response.put(ResultMapper.DATA, arr);
        log.info("[count]{}", sr.getHits().getTotalHits());
        return response.toString();
    }


    /**
     * 高亮显示
     *
     * @param param
     */
    private static void buildHighlight(QueryParam param, SearchSourceBuilder searchSourceBuilder) {
        HighlightBuilder highlighter = new HighlightBuilder();
        highlighter.preTags("<em style='color:#1772fa;font-style:normal;'>");
        highlighter.postTags("</em>");
        String[] highlight = param.getHighlight();//高亮字段
        if (ArrayUtil.isNotEmpty(highlight)) {
            for (String field : highlight) {
                highlighter.field(field);
            }
            highlighter.fragmentSize(Integer.MAX_VALUE);
            searchSourceBuilder.highlighter(highlighter);
        }
    }

    /**
     * 排序
     *
     * @param param
     * @param searchSourceBuilder
     * @return
     */
    private static void buildSort(QueryParam param, SearchSourceBuilder searchSourceBuilder) {
        Map<String, String> sortMap = param.getSortMap();//排序
        if (sortMap != null) {
            for (String sortKey : sortMap.keySet()) {
                if (sortKey.equals("_score")) {
                    searchSourceBuilder.sort(new ScoreSortBuilder().order(SortOrder.fromString(sortMap.get(sortKey))));
                } else {
                    searchSourceBuilder.sort(SortBuilders.fieldSort(sortKey).order(SortOrder.fromString(sortMap.get(sortKey))));
                }
            }
        }
    }

    /**
     * 重打分
     *
     * @param param
     * @param searchSourceBuilder
     * @return
     */
    private static void buildRescore(QueryParam param, SearchSourceBuilder searchSourceBuilder) {
        QueryBuilder rescoreQuery = param.getRescoreQuery();
        if (rescoreQuery != null) {
            QueryRescorerBuilder builder = new QueryRescorerBuilder(rescoreQuery);
            searchSourceBuilder.addRescorer(builder
                    .windowSize(100)
                    .setQueryWeight(0.1F)
                    .setRescoreQueryWeight(1F)
                    .setScoreMode(QueryRescoreMode.Total)
            );
        }
    }


    /**
     * 返回字段设置
     *
     * @param param
     * @param searchSourceBuilder
     * @return
     */
    private static void buildFetchSource(QueryParam param, SearchSourceBuilder searchSourceBuilder) {
        Map<String, String> excludeMap = param.getExcludeMap();
        Map<String, String> includeMap = param.getIncludeMap();
        String[] excludes = new String[]{};
        String[] includes = new String[]{};
        if (MapUtils.isNotEmpty(excludeMap)) {
            excludes = excludeMap.entrySet().stream().map(entry -> entry.getKey()).toArray(String[]::new);
        }
        if (MapUtils.isNotEmpty(includeMap)) {
            includes = includeMap.entrySet().stream().map(entry -> entry.getKey()).toArray(String[]::new);
        }
        if (ArrayUtil.isNotEmpty(excludes) || ArrayUtil.isNotEmpty(excludes)) {
            searchSourceBuilder.fetchSource(includes, excludes);
        }
    }

}
